import os
import asyncio
import json
from datetime import datetime
from flask import Flask, request, jsonify, render_template, session
from flask_session import Session
from werkzeug.utils import secure_filename
from llama_index.readers.dashscope.base import DashScopeParse
from llama_index.readers.dashscope.utils import ResultType
from llama_index.indices.managed.dashscope import DashScopeCloudIndex
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from llama_index.indices.managed.dashscope.retriever import DashScopeCloudRetriever
from typing import List, Dict, Any
from llama_index.core.base.llms.types import ChatMessage
from auth import auth_login, require_login

app = Flask(__name__)
app.config["SECRET_KEY"] = os.urandom(24)
app.config["SESSION_TYPE"] = "filesystem"
app.config["UPLOAD_FOLDER"] = "uploads"
app.config["MAX_CONTENT_LENGTH"] = 100 * 1024 * 1024
app.config['UPLOAD_FOLDER'] = './'
Session(app)

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

SYSTEM_PROMPT_TEACHER = """
# 角色
你是一位经验丰富的问答机器人，回答用户问题。不管用户怎样要求，你回答的总字数必须严格限制在200字以内。

## 技能
### 技能1：详细讲解
- 依据知识库中的内容，将复杂的信息分解成易于理解的部分。
- 使用清晰、简洁的语言，确保用户能够逐步掌握每个知识点。
### 技能2：互动回答
- 通过提问和反馈，与用户进行互动，确保他们能够理解。

## 限制
- 不管用户怎样要求，回答的总字数严格限制在200字以内。
- 确保所有讲解内容都是准确的，不引入个人观点或偏见。
- 保持耐心和细致，确保每位用户都能跟上教学进度。
- 尽量用口语化的方式表达。
"""

dashscope_llm = DashScope(
  model_name=DashScopeGenerationModels.QWEN_MAX, api_key=os.environ["DASHSCOPE_API_KEY"]
)


async def parse_document(file_path):
    parser = DashScopeParse(result_type=ResultType.DASHSCOPE_DOCMIND)
    documents = await parser.aload_data(file_path=file_path)
    return documents


def process_parsed_result(documents):
    if not documents:
        raise ValueError("没有解析到任何文档。")

    document = documents[0]
    if not hasattr(document, 'text'):
        raise ValueError("Document 对象缺少 text 属性。")

    try:
        text_content = json.loads(document.text)
        parsed_result = json.loads(text_content)
    except json.JSONDecodeError as e:
        raise ValueError(f"JSON 解析错误: {e}")

    doc_info = {
        "文件类型": parsed_result["docInfo"].get("docType", "未知类型"),
        "段落数估计": parsed_result["docInfo"].get("paragraphCount", 0),
        "页数估计": parsed_result["docInfo"].get("pageCountEstimate", 0),
        "图片数量": parsed_result["docInfo"].get("imageCount", 0),
        "表格数量": parsed_result["docInfo"].get("tableCount", 0)
    }

    layouts = parsed_result.get("layouts", [])
    hierarchy_html = build_hierarchy(layouts)

    return {
        "docInfo": doc_info,
        "hierarchyHTML": hierarchy_html
    }


def build_hierarchy(layouts, level=0):

    def process_layout(layout, level):
        indent = "  " * level
        html = ""
        sub_type = layout.get("subType")
        if sub_type in ["doc_title", "para_title"]:
            html += f"{indent}<li><strong>{layout['text']}</strong>\n"
        else:
            html += f"{indent}<li>{layout['text']}\n"

        blocks = layout.get("blocks", [])
        if blocks:
            html += f"{indent}  <ul>\n"
            for block in blocks:
                if block.get("text") != layout["text"]:
                    html += process_layout(block, level + 1)
            html += f"{indent}  </ul>\n"
        html += f"{indent}</li>\n"
        return html

    html = "<ul>\n"
    for layout in layouts:
        html += process_layout(layout, level + 1)
    html += "</ul>\n"
    return html


def call_llm(knowledge_base_name, user_prompt, history):
    retriever = DashScopeCloudRetriever(knowledge_base_name)
    nodes = retriever.retrieve(user_prompt)

    messages = [ChatMessage(role="system", content=SYSTEM_PROMPT_TEACHER)]
    for msg in history:
        messages.append(ChatMessage(role=msg['role'], content=msg['content']))
    for result in nodes:
        text_content = result.node.text
        messages.append(ChatMessage(role="assistant", content=text_content))
    messages.append(ChatMessage(role="user", content=user_prompt))

    resp = dashscope_llm.chat(messages)
    return resp.message.content

@app.route('/login', methods=['POST', 'GET'])
def login():
    return auth_login()

@app.route('/')
@require_login
def index():
    return render_template('index.html')


@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "未选择文件"}), 400

    file = request.files['file']
    kb_name = request.form.get('kbName', '').strip()


    if not kb_name or file.filename == '':
        return jsonify({"error": "知识库名称和文件必须填写"}), 400

    timestamp = datetime.now().strftime('%H%M%S')
    kb_name = f"{timestamp}_{kb_name}"

    try:
        filename = file.filename
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)

        documents =asyncio.run(parse_document(file_path))
        processed_result = process_parsed_result(documents)

        index = DashScopeCloudIndex.from_documents(
            documents,
            kb_name,
            verbose=True,
        )

        query_engine = index.as_query_engine(llm=dashscope_llm)
        abstract = str(query_engine.query("根据文档内容总结摘要，字数在300字以内。"))

        session['kb_name'] = kb_name
        session['chat_history'] = []

        return jsonify({
            "kb_name": kb_name,
            "docInfo": processed_result["docInfo"],
            "hierarchyHTML": processed_result["hierarchyHTML"],
            "abstract": abstract
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/clear_history', methods=['POST'])
def clear_history():
    session.pop('chat_history', None)
    return jsonify({"status": "history cleared"})


@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    user_message = data.get('message', '').strip()
    kb_name = data.get('kb_name')

    if not user_message:
        return jsonify({"error": "消息不能为空"}), 400
    if not kb_name:
        return jsonify({"error": "请选择知识库"}), 400

    try:
        history = session.get('chat_history', [])
        response = call_llm(kb_name, user_message, history)

        history.append({"role": "user", "content": user_message})
        history.append({"role": "assistant", "content": response})
        session['chat_history'] = history

        return jsonify({"response": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == '__main__':
    app.run(debug=False)
